package com.dmall.sort.com.dmall.heap;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

public class MaxHeap<E extends Comparable<E>> implements Heap<E> {

    private List<E> data;

    public MaxHeap(int capacity) {
        data = new ArrayList<>(capacity);
    }

    public MaxHeap() {
        this(10);
    }

    public MaxHeap(E[] arr) {
        data = Arrays.asList(arr);
        for (int i = parent(data.size() - 1); i >= 0; i--) {
            shiftDown(i);
        }
    }

    @Override
    public int size() {
        return data.size();
    }

    @Override
    public boolean isEmpty() {
        return data.isEmpty();
    }

    @Override
    public void add(E e) {
        data.add(e);
        shiftUp(data.size() - 1);
    }

    @Override
    public E extract() {
        swap(0, data.size() - 1);
        E result = data.remove(data.size() - 1);
        shiftDown(0);
        return result;
    }

    @Override
    public E peek() {
        if (data.isEmpty()) {
            throw new IllegalArgumentException("MaxHeap is empty.");
        }
        return data.get(0);
    }

    private void shiftUp(int index) {
        while (index > 0 && data.get(index).compareTo(data.get(parent(index))) > 0) {
            swap(index, parent(index));
            index = parent(index);
        }
    }

    private void shiftDown(int index) {
        while (left(index) < data.size()) {
            int maxIndex = left(index);
            if (right(index) < data.size() && data.get(left(index)).compareTo(data.get(right(index))) < 0) {
                maxIndex ++;
            }
            if (data.get(index).compareTo(data.get(maxIndex)) >= 0) {
                break;
            }
            swap(index, maxIndex);
            index = maxIndex;
        }
    }

    private int parent(int index) {
        if (index == 0) {
            throw new IllegalArgumentException("index-0 doesn't have parent.");
        }
        return (index - 1) / 2;
    }

    private int left(int index) {
        return index * 2 + 1;
    }

    private int right(int index) {
        return index * 2 + 2;
    }

    private void swap(int i, int j) {
        E temp = data.get(i);
        data.set(i, data.get(j));
        data.set(j, temp);
    }
}
